from matplotlib import pyplot
import math

def RshiftedLJ(eps, sigma, r0, r):
    rc=sigma+2.5
    if r<=r0:
        fun=1000000.0
    elif r>r0 and r<rc:
        fun=4.0*eps*(math.pow(sigma/(r-r0),12.0)-math.pow(sigma/(r-r0),6.0))
    else:
        fun=0.0
    return fun

def magnetica(u0, theta, r):
    fun=-u0*math.pow(r,-3.0)*(1.0-3.0*math.pow(math.cos(theta),2.0))
    return fun


sigma=0.1
r0=0.9
eps=1.0
theta=math.pi
u0=20.0

desde=0.99
hasta=1.2
puntos=100
paso=(hasta-desde)/puntos

x=[]
y=[]
y2=[]
y3=[]
y4=[]
for i in range(puntos):
    r=desde+i*paso
    x.append(r)
    y.append(RshiftedLJ(0.01, sigma, r0, r))
    y2.append(RshiftedLJ(0.1, sigma, r0, r))
    y3.append(RshiftedLJ(0.5, sigma, r0, r))
    y4.append(RshiftedLJ(1.0, sigma, r0, r))
    #y2.append(magnetica(u0, theta, r))
    #y3.append(RshiftedLJ(eps, sigma, r0, r)+magnetica(u0, theta, r))

p1,=pyplot.plot(x,y,'ro-')
p2,=pyplot.plot(x,y2,'go-')
p3,=pyplot.plot(x,y3,'bo-')
p4,=pyplot.plot(x,y4,'bo-')
pyplot.legend([p1,p2,p3,p4],['0.01','0.1','0.5','1.0'])
#pyplot.legend([p2,p3],["Interaccion magnetica","R-Shifted + magnetica"])
pyplot.xlabel("Distancia (diametros)")
pyplot.ylabel("Energia/kBT")
pyplot.show()